/*
 * Copyright 2019 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#ifndef SkRuntimeEffect_DEFINED
#define SkRuntimeEffect_DEFINED

#include "include/core/SkBlender.h"     // IWYU pragma: keep
#include "include/core/SkColorFilter.h" // IWYU pragma: keep
#include "include/core/SkData.h"
#include "include/core/SkFlattenable.h"
#include "include/core/SkMatrix.h"
#include "include/core/SkRefCnt.h"
#include "include/core/SkShader.h"
#include "include/core/SkSpan.h"
#include "include/core/SkString.h"
#include "include/core/SkTypes.h"
#include "include/private/SkSLSampleUsage.h"
#include "include/private/base/SkOnce.h"
#include "include/private/base/SkTemplates.h"
#include "include/private/base/SkTo.h"
#include "include/private/base/SkTypeTraits.h"
#include "include/sksl/SkSLDebugTrace.h"
#include "include/sksl/SkSLVersion.h"

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

struct SkIPoint;

namespace SkSL {
class DebugTracePriv;
class FunctionDefinition;
struct Program;
enum class ProgramKind : int8_t;
struct ProgramSettings;
} // namespace SkSL

namespace SkSL::RP {
class Program;
}

/*
 * SkRuntimeEffect supports creating custom SkShader and SkColorFilter objects using Skia's SkSL
 * shading language.
 *
 * NOTE: This API is experimental and subject to change.
 */
class SK_API SkRuntimeEffect : public SkRefCnt {
public:
    // Reflected description of a uniform variable in the effect's SkSL
    struct SK_API Uniform {
        enum class Type {
            kFloat,
            kFloat2,
            kFloat3,
            kFloat4,
            kFloat2x2,
            kFloat3x3,
            kFloat4x4,
            kInt,
            kInt2,
            kInt3,
            kInt4,
        };

        enum Flags {
            // Uniform is declared as an array. 'count' contains array length.
            kArray_Flag = 0x1,

            // Uniform is declared with layout(color). Colors should be supplied as unpremultiplied,
            // extended-range (unclamped) sRGB (ie SkColor4f). The uniform will be automatically
            // transformed to unpremultiplied extended-range working-space colors.
            kColor_Flag = 0x2,

            // When used with SkMeshSpecification, indicates that the uniform is present in the
            // vertex shader. Not used with SkRuntimeEffect.
            kVertex_Flag = 0x4,

            // When used with SkMeshSpecification, indicates that the uniform is present in the
            // fragment shader. Not used with SkRuntimeEffect.
            kFragment_Flag = 0x8,

            // This flag indicates that the SkSL uniform uses a medium-precision type
            // (i.e., `half` instead of `float`).
            kHalfPrecision_Flag = 0x10,
        };

        std::string_view name;
        size_t offset;
        Type type;
        int count;
        uint32_t flags;

        bool isArray() const
        {
            return SkToBool(this->flags & kArray_Flag);
        }
        bool isColor() const
        {
            return SkToBool(this->flags & kColor_Flag);
        }
        size_t sizeInBytes() const;
    };

    // Reflected description of a uniform child (shader or colorFilter) in the effect's SkSL
    enum class ChildType {
        kShader,
        kColorFilter,
        kBlender,
    };

    struct Child {
        std::string_view name;
        ChildType type;
        int index;
    };

    class Options {
    public:
        // For testing purposes, disables optimization and inlining. (Normally, Runtime Effects
        // don't run the inliner directly, but they still get an inlining pass once they are
        // painted.)
        bool forceUnoptimized = false;

    private:
        friend class SkRuntimeEffect;
        friend class SkRuntimeEffectPriv;

        // This flag allows Runtime Effects to access Skia implementation details like sk_FragCoord
        // and functions with private identifiers (e.g. $rgb_to_hsl).
        bool allowPrivateAccess = false;

        // TODO(skia:11209) - Replace this with a promised SkCapabilities?
        // This flag lifts the ES2 restrictions on Runtime Effects that are gated by the
        // `strictES2Mode` check. Be aware that the software renderer and pipeline-stage effect are
        // still largely ES3-unaware and can still fail or crash if post-ES2 features are used.
        // This is only intended for use by tests and certain internally created effects.
        SkSL::Version maxVersionAllowed = SkSL::Version::k100;
    };

    // If the effect is compiled successfully, `effect` will be non-null.
    // Otherwise, `errorText` will contain the reason for failure.
    struct Result {
        sk_sp<SkRuntimeEffect> effect;
        SkString errorText;
    };

    // MakeForColorFilter and MakeForShader verify that the SkSL code is valid for those stages of
    // the Skia pipeline. In all of the signatures described below, color parameters and return
    // values are flexible. They are listed as being 'vec4', but they can also be 'half4' or
    // 'float4'. ('vec4' is an alias for 'float4').

    // We can't use a default argument for `options` due to a bug in Clang.
    // https://bugs.llvm.org/show_bug.cgi?id=36684

    // Color filter SkSL requires an entry point that looks like:
    //     vec4 main(vec4 inColor) { ... }
    static Result MakeForColorFilter(SkString sksl, const Options &);
    static Result MakeForColorFilter(SkString sksl)
    {
        return MakeForColorFilter(std::move(sksl), Options{});
    }

    // Shader SkSL requires an entry point that looks like:
    //     vec4 main(vec2 inCoords) { ... }
    static Result MakeForShader(SkString sksl, const Options &);
    static Result MakeForShader(SkString sksl)
    {
        return MakeForShader(std::move(sksl), Options{});
    }

    // Blend SkSL requires an entry point that looks like:
    //     vec4 main(vec4 srcColor, vec4 dstColor) { ... }
    static Result MakeForBlender(SkString sksl, const Options &);
    static Result MakeForBlender(SkString sksl)
    {
        return MakeForBlender(std::move(sksl), Options{});
    }

    // Object that allows passing a SkShader, SkColorFilter or SkBlender as a child
    class SK_API ChildPtr {
    public:
        ChildPtr() = default;
        ChildPtr(sk_sp<SkShader> s) : fChild(std::move(s)) {}
        ChildPtr(sk_sp<SkColorFilter> cf) : fChild(std::move(cf)) {}
        ChildPtr(sk_sp<SkBlender> b) : fChild(std::move(b)) {}

        // Asserts that the flattenable is either null, or one of the legal derived types
        ChildPtr(sk_sp<SkFlattenable> f);

        std::optional<ChildType> type() const;

        SkShader *shader() const;
        SkColorFilter *colorFilter() const;
        SkBlender *blender() const;
        SkFlattenable *flattenable() const
        {
            return fChild.get();
        }

        using sk_is_trivially_relocatable = std::true_type;

    private:
        sk_sp<SkFlattenable> fChild;

        static_assert(::sk_is_trivially_relocatable<decltype(fChild)>::value);
    };

    sk_sp<SkShader> makeShader(sk_sp<const SkData> uniforms, sk_sp<SkShader> children[], size_t childCount,
        const SkMatrix *localMatrix = nullptr) const;
    sk_sp<SkShader> makeShader(sk_sp<const SkData> uniforms, SkSpan<const ChildPtr> children,
        const SkMatrix *localMatrix = nullptr) const;

    sk_sp<SkColorFilter> makeColorFilter(sk_sp<const SkData> uniforms) const;
    sk_sp<SkColorFilter> makeColorFilter(sk_sp<const SkData> uniforms, sk_sp<SkColorFilter> children[],
        size_t childCount) const;
    sk_sp<SkColorFilter> makeColorFilter(sk_sp<const SkData> uniforms, SkSpan<const ChildPtr> children) const;

    sk_sp<SkBlender> makeBlender(sk_sp<const SkData> uniforms, SkSpan<const ChildPtr> children = {}) const;

    /* *
     * Creates a new Runtime Effect patterned after an already-existing one. The new shader behaves
     * like the original, but also creates a debug trace of its execution at the requested
     * coordinate. After painting with this shader, the associated DebugTrace object will contain a
     * shader execution trace. Call `writeTrace` on the debug trace object to generate a full trace
     * suitable for a debugger, or call `dump` to emit a human-readable trace.
     *
     * Debug traces are only supported on a raster (non-GPU) canvas.

     * Debug traces are currently only supported on shaders. Color filter and blender tracing is a
     * work-in-progress.
     */
    struct TracedShader {
        sk_sp<SkShader> shader;
        sk_sp<SkSL::DebugTrace> debugTrace;
    };
    static TracedShader MakeTraced(sk_sp<SkShader> shader, const SkIPoint &traceCoord);

    // Returns the SkSL source of the runtime effect shader.
    const std::string &source() const;

    // Combined size of all 'uniform' variables. When calling makeColorFilter or makeShader,
    // provide an SkData of this size, containing values for all of those variables.
    size_t uniformSize() const;

    SkSpan<const Uniform> uniforms() const
    {
        return SkSpan(fUniforms);
    }
    SkSpan<const Child> children() const
    {
        return SkSpan(fChildren);
    }

    // Returns pointer to the named uniform variable's description, or nullptr if not found
    const Uniform *findUniform(std::string_view name) const;

    // Returns pointer to the named child's description, or nullptr if not found
    const Child *findChild(std::string_view name) const;

    // Allows the runtime effect type to be identified.
    bool allowShader() const
    {
        return (fFlags & kAllowShader_Flag);
    }
    bool allowColorFilter() const
    {
        return (fFlags & kAllowColorFilter_Flag);
    }
    bool allowBlender() const
    {
        return (fFlags & kAllowBlender_Flag);
    }

    static void RegisterFlattenables();
    ~SkRuntimeEffect() override;

private:
    enum Flags {
        kUsesSampleCoords_Flag = 0x001,
        kAllowColorFilter_Flag = 0x002,
        kAllowShader_Flag = 0x004,
        kAllowBlender_Flag = 0x008,
        kSamplesOutsideMain_Flag = 0x010,
        kUsesColorTransform_Flag = 0x020,
        kAlwaysOpaque_Flag = 0x040,
        kAlphaUnchanged_Flag = 0x080,
        kDisableOptimization_Flag = 0x100,
    };

    SkRuntimeEffect(std::unique_ptr<SkSL::Program> baseProgram, const Options &options,
        const SkSL::FunctionDefinition &main, std::vector<Uniform> &&uniforms, std::vector<Child> &&children,
        std::vector<SkSL::SampleUsage> &&sampleUsages, uint32_t flags);

    sk_sp<SkRuntimeEffect> makeUnoptimizedClone();

    static Result MakeFromSource(SkString sksl, const Options &options, SkSL::ProgramKind kind);

    static Result MakeInternal(std::unique_ptr<SkSL::Program> program, const Options &options, SkSL::ProgramKind kind);

    static SkSL::ProgramSettings MakeSettings(const Options &options);

    uint32_t hash() const
    {
        return fHash;
    }
    bool usesSampleCoords() const
    {
        return (fFlags & kUsesSampleCoords_Flag);
    }
    bool samplesOutsideMain() const
    {
        return (fFlags & kSamplesOutsideMain_Flag);
    }
    bool usesColorTransform() const
    {
        return (fFlags & kUsesColorTransform_Flag);
    }
    bool alwaysOpaque() const
    {
        return (fFlags & kAlwaysOpaque_Flag);
    }
    bool isAlphaUnchanged() const
    {
        return (fFlags & kAlphaUnchanged_Flag);
    }

    const SkSL::RP::Program *getRPProgram(SkSL::DebugTracePriv *debugTrace) const;

    friend class GrSkSLFP;             // usesColorTransform
    friend class SkRuntimeShader;      // fBaseProgram, fMain, fSampleUsages, getRPProgram()
    friend class SkRuntimeBlender;     //
    friend class SkRuntimeColorFilter; //

    friend class SkRuntimeEffectPriv;

    uint32_t fHash;

    std::unique_ptr<SkSL::Program> fBaseProgram;
    std::unique_ptr<SkSL::RP::Program> fRPProgram;
    mutable SkOnce fCompileRPProgramOnce;
    const SkSL::FunctionDefinition &fMain;
    std::vector<Uniform> fUniforms;
    std::vector<Child> fChildren;
    std::vector<SkSL::SampleUsage> fSampleUsages;

    uint32_t fFlags; // Flags
};

/* * Base class for SkRuntimeShaderBuilder, defined below. */
class SkRuntimeEffectBuilder {
public:
    struct BuilderUniform {
        // Copy 'val' to this variable. No type conversion is performed - 'val' must be same
        // size as expected by the effect. Information about the variable can be queried by
        // looking at fVar. If the size is incorrect, no copy will be performed, and debug
        // builds will abort. If this is the result of querying a missing variable, fVar will
        // be nullptr, and assigning will also do nothing (and abort in debug builds).
        template <typename T>
        std::enable_if_t<std::is_trivially_copyable<T>::value, BuilderUniform &> operator = (const T &val)
        {
            if (!fVar) {
                SkDEBUGFAIL("Assigning to missing variable");
            } else if (sizeof(val) != fVar->sizeInBytes()) {
                SkDEBUGFAIL("Incorrect value size");
            } else {
                memcpy(SkTAddOffset<void>(fOwner->writableUniformData(), fVar->offset), &val, sizeof(val));
            }
            return *this;
        }

        BuilderUniform &operator = (const SkMatrix &val)
        {
            if (!fVar) {
                SkDEBUGFAIL("Assigning to missing variable");
            } else if (fVar->sizeInBytes() != 9 * sizeof(float)) {
                SkDEBUGFAIL("Incorrect value size");
            } else {
                float *data = SkTAddOffset<float>(fOwner->writableUniformData(), (ptrdiff_t)fVar->offset);
                data[0] = val.get(0);
                data[1] = val.get(3);
                data[2] = val.get(6);
                data[3] = val.get(1);
                data[4] = val.get(4);
                data[5] = val.get(7);
                data[6] = val.get(2);
                data[7] = val.get(5);
                data[8] = val.get(8);
            }
            return *this;
        }

        template <typename T> bool set(const T val[], const int count)
        {
            static_assert(std::is_trivially_copyable<T>::value, "Value must be trivial copyable");
            if (!fVar) {
                SkDEBUGFAIL("Assigning to missing variable");
                return false;
            } else if (sizeof(T) * count != fVar->sizeInBytes()) {
                SkDEBUGFAIL("Incorrect value size");
                return false;
            } else {
                memcpy(SkTAddOffset<void>(fOwner->writableUniformData(), fVar->offset), val, sizeof(T) * count);
            }
            return true;
        }

        SkRuntimeEffectBuilder *fOwner;
        const SkRuntimeEffect::Uniform *fVar; // nullptr if the variable was not found
    };

    struct BuilderChild {
        template <typename T> BuilderChild &operator = (sk_sp<T> val)
        {
            if (!fChild) {
                SkDEBUGFAIL("Assigning to missing child");
            } else {
                fOwner->fChildren[(size_t)fChild->index] = std::move(val);
            }
            return *this;
        }

        BuilderChild &operator = (std::nullptr_t)
        {
            if (!fChild) {
                SkDEBUGFAIL("Assigning to missing child");
            } else {
                fOwner->fChildren[(size_t)fChild->index] = SkRuntimeEffect::ChildPtr{};
            }
            return *this;
        }

        SkRuntimeEffectBuilder *fOwner;
        const SkRuntimeEffect::Child *fChild; // nullptr if the child was not found
    };

    const SkRuntimeEffect *effect() const
    {
        return fEffect.get();
    }

    BuilderUniform uniform(std::string_view name)
    {
        return { this, fEffect->findUniform(name) };
    }
    BuilderChild child(std::string_view name)
    {
        return { this, fEffect->findChild(name) };
    }

    // Get access to the collated uniforms and children (in the order expected by APIs like
    // makeShader on the effect):
    sk_sp<const SkData> uniforms() const
    {
        return fUniforms;
    }
    SkSpan<const SkRuntimeEffect::ChildPtr> children() const
    {
        return fChildren;
    }

protected:
    SkRuntimeEffectBuilder() = delete;
    explicit SkRuntimeEffectBuilder(sk_sp<SkRuntimeEffect> effect)
        : fEffect(std::move(effect)),
          fUniforms(SkData::MakeZeroInitialized(fEffect->uniformSize())),
          fChildren(fEffect->children().size())
    {}
    explicit SkRuntimeEffectBuilder(sk_sp<SkRuntimeEffect> effect, sk_sp<SkData> uniforms)
        : fEffect(std::move(effect)), fUniforms(std::move(uniforms)), fChildren(fEffect->children().size())
    {}

    SkRuntimeEffectBuilder(SkRuntimeEffectBuilder &&) = default;
    SkRuntimeEffectBuilder(const SkRuntimeEffectBuilder &) = default;

    SkRuntimeEffectBuilder &operator = (SkRuntimeEffectBuilder &&) = delete;
    SkRuntimeEffectBuilder &operator = (const SkRuntimeEffectBuilder &) = delete;

private:
    void *writableUniformData()
    {
        if (!fUniforms->unique()) {
            fUniforms = SkData::MakeWithCopy(fUniforms->data(), fUniforms->size());
        }
        return fUniforms->writable_data();
    }

    sk_sp<SkRuntimeEffect> fEffect;
    sk_sp<SkData> fUniforms;
    std::vector<SkRuntimeEffect::ChildPtr> fChildren;
};

/* *
 * SkRuntimeShaderBuilder is a utility to simplify creating SkShader objects from SkRuntimeEffects.
 *
 * NOTE: Like SkRuntimeEffect, this API is experimental and subject to change!
 *
 * Given an SkRuntimeEffect, the SkRuntimeShaderBuilder manages creating an input data block and
 * provides named access to the 'uniform' variables in that block, as well as named access
 * to a list of child shader slots. Usage:
 *
 * sk_sp<SkRuntimeEffect> effect = ...;
 * SkRuntimeShaderBuilder builder(effect);
 * builder.uniform("some_uniform_float")  = 3.14f;
 * builder.uniform("some_uniform_matrix") = SkM44::Rotate(...);
 * builder.child("some_child_effect")     = mySkImage->makeShader(...);
 * ...
 * sk_sp<SkShader> shader = builder.makeShader(nullptr, false);
 *
 * Note that SkRuntimeShaderBuilder is built entirely on the public API of SkRuntimeEffect,
 * so can be used as-is or serve as inspiration for other interfaces or binding techniques.
 */
class SK_API SkRuntimeShaderBuilder : public SkRuntimeEffectBuilder {
public:
    explicit SkRuntimeShaderBuilder(sk_sp<SkRuntimeEffect>);
    // This is currently required by Android Framework but may go away if that dependency
    // can be removed.
    SkRuntimeShaderBuilder(const SkRuntimeShaderBuilder &) = default;
    ~SkRuntimeShaderBuilder();

    sk_sp<SkShader> makeShader(const SkMatrix *localMatrix = nullptr) const;

private:
    explicit SkRuntimeShaderBuilder(sk_sp<SkRuntimeEffect> effect, sk_sp<SkData> uniforms)
        : SkRuntimeEffectBuilder(std::move(effect), std::move(uniforms))
    {}

    friend class SkRuntimeImageFilter;
};

/* *
 * SkRuntimeColorFilterBuilder makes it easy to setup and assign uniforms to runtime color filters.
 */
class SK_API SkRuntimeColorFilterBuilder : public SkRuntimeEffectBuilder {
public:
    explicit SkRuntimeColorFilterBuilder(sk_sp<SkRuntimeEffect>);
    ~SkRuntimeColorFilterBuilder();

    SkRuntimeColorFilterBuilder(const SkRuntimeColorFilterBuilder &) = delete;
    SkRuntimeColorFilterBuilder &operator = (const SkRuntimeColorFilterBuilder &) = delete;

    sk_sp<SkColorFilter> makeColorFilter() const;
};

/* *
 * SkRuntimeBlendBuilder is a utility to simplify creation and uniform setup of runtime blenders.
 */
class SK_API SkRuntimeBlendBuilder : public SkRuntimeEffectBuilder {
public:
    explicit SkRuntimeBlendBuilder(sk_sp<SkRuntimeEffect>);
    ~SkRuntimeBlendBuilder();

    SkRuntimeBlendBuilder(const SkRuntimeBlendBuilder &) = delete;
    SkRuntimeBlendBuilder &operator = (const SkRuntimeBlendBuilder &) = delete;

    sk_sp<SkBlender> makeBlender() const;
};

#endif // SkRuntimeEffect_DEFINED
